import os.path as osp
import pandas as pd
import torch
import torch.nn.functional as F
import torch.nn as nn
import random
import numpy as np
import pylab
from PIL import Image
from torch_geometric.data import Dataset, download_url, Data
from torchvision import transforms
import pdb

def GasussianKernel(size=41):
    if size%2 == 0:
        size = size + 1 

    cx = size//2 
    cy = size//2
    matrix = torch.zeros((size,size),dtype=torch.float)
    for i in range(0,size):
        for j in range(0,size):
            matrix[i,j] = np.exp(-np.abs(i-cx)/10-np.abs(j-cy)/10)
    return matrix

class GaussianBlur(nn.Module):
    def __init__(self,size=41):
        super(GaussianBlur,self).__init__()
        self.size = size
        kernel = GasussianKernel(size).unsqueeze(0).unsqueeze(0)
        self.weight = nn.Parameter(data=kernel,requires_grad=False)
    def forward(self,x):
        x = F.conv2d(x.unsqueeze(0).unsqueeze(0),self.weight,padding=self.size//2)
        return x


class ISPD15Dataset(Dataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(ISPD15Dataset, self).__init__(root, transform, pre_transform)

    @property
    def raw_file_names(self):
        return ['des_perf_a', 'edit_dist_a', 'fft_a','fft_b','matrix_mult_a','matrix_mult_b','matrix_mult_c','pci_bridge32_a','pci_bridge32_b']

    @property
    def processed_file_names(self):
        return ['data_%d.pt'%i for i in range(0,9)]

    def process(self):
        i = 0
        for design in self.raw_file_names:
            path = osp.join(self.raw_dir,design)
            net_path = osp.join(path,'nets.txt')
            nodes_path = osp.join(path,'node_feature.txt')
            name_path = osp.join(path,'names.txt')
            pos_path = osp.join(path,'pos.txt')
            graph_lab_path = osp.join(path,'graph_labels.txt')
            pin_path = osp.join(path,'pins.txt')
            region_path = osp.join(path,'region.txt')
            # Read data from `raw_path`.
            nets = np.array(pd.read_table(net_path,header=None))
            node_lab = np.array(pd.read_table(nodes_path,header=None))
            pin_feature = np.array(pd.read_table(pin_path,header=None))[:,2:]
            xl,yl,xh,yh = np.loadtxt(region_path)
            # normalize
            node_lab[:,0] /= xh-xl
            node_lab[:,1] /= yh-yl
            pin_feature[:,0] /= xh-xl
            pin_feature[:,1] /= yh-yl

            x = torch.tensor(node_lab, dtype=torch.float)
            edge_index = torch.tensor(nets.T, dtype=torch.long)
            pin_feature = torch.tensor(pin_feature,dtype=torch.float)
            data = Data(x=x, edge_index=edge_index,y=random.random(),pin_feature=pin_feature)

            if self.pre_filter is not None and not self.pre_filter(data):
                continue

            if self.pre_transform is not None:
                data = self.pre_transform(data)

            torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
            i += 1

    def len(self):
        return len(self.processed_file_names)

    def get(self, idx):
        data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(idx)))
        return data


class MyDataset(Dataset):
    def __init__(self, root, transform=None, pre_transform=None, classes=20):
        self.tot_file_num = None # int
        self.file_num = None # dict, file nums for each design
        self.ptr = None
        self.classes = classes
        self.y_label_type = None
        super(MyDataset, self).__init__(root, transform, pre_transform)

    @property
    def raw_file_names(self):
        return ['des_perf_a', 'edit_dist_a', 'fft_a','fft_b','matrix_mult_a','matrix_mult_b','matrix_mult_c','pci_bridge32_b']

    @property
    def num_features(self):
        return self[0].x.size(1) + self[0].pin_feature.size(1)  + 1

    @property
    def num_classes(self):
        return int(self.classes)

    @property
    def processed_file_names(self):
        if self.tot_file_num is None:
            self.tot_file_num = 0
            self.file_num = {}
            self.ptr = {}
            for design in self.raw_file_names:
                path = osp.join(self.raw_dir,design)
                name_path = osp.join(path,'names.txt')
                names = np.array(pd.read_table(name_path,header=None)).reshape(-1)
                self.tot_file_num += names.shape[0]
                self.file_num[design] = names.shape[0]
            self.ptr[self.raw_file_names[0]] = 0
            for i in range(1,int(len(self.raw_file_names))):
                self.ptr[self.raw_file_names[i]] = self.ptr[self.raw_file_names[i-1]] + self.file_num[self.raw_file_names[i-1]]
        return ['data_%d.pt'%i for i in range(0,self.tot_file_num)]



    def reset_y(self,classes=None):
        if classes is not None:
            self.classes = classes
        mylist = np.array([])
        for design in self.raw_file_names:
            # paths
            path = osp.join(self.raw_dir,design)
            name_path = osp.join(path,'names.txt')
            graph_lab_path = osp.join(path,'graph_labels.txt')
            region_path = osp.join(path,'region.txt')
            # Read data from `raw_path`.
            xl,yl,xh,yh = np.loadtxt(region_path)
            with open(graph_lab_path,'r') as f:
                rWLs = np.array([float(line) if line != '\n' else 0 for line in f.readlines()])
            with open(name_path,'r') as f:
                names = [int(line) if line != '\n' else 0 for line in f.readlines()]
            # normalize
            rWLs = rWLs/(xh-xl+yh-yl)
            rWLs = rWLs/1000.0

            mylist = np.append(mylist,rWLs[names])
        
        # stastics
        mylist = np.array(mylist).reshape(-1,)
        rel = mylist.tolist()
        pylab.hist(rel,20)
        pylab.xlabel('Range')
        pylab.ylabel('Count')
        pylab.savefig('{}.png'.format('Total'))
        pylab.cla()


        maxx = int(mylist.max())+1
        step = maxx/self.classes
        for i in range(0,self.tot_file_num):
            data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
            data.y = torch.tensor(int(data.wl.item() // step),dtype=torch.long)
            if data.y >= self.classes:
                data.y = self.classes - 1
            torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))

    def set_y_as_wl(self):
        if self.y_label_type == 'wl':
            return
        self.y_label_type = 'wl'
        for i in range(0,self.tot_file_num):
            data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
            data.y = torch.tensor(data.wl,dtype=torch.float)
            torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))


    def set_y_as_rank(self,classes=None):
        if self.y_label_type == 'rank' and self.classes == classes:
            return
        self.y_label_type == 'rank'
        if classes is not None:
            self.classes = classes
        mylist = np.array([])
        for design in self.raw_file_names:
            # paths
            path = osp.join(self.raw_dir,design)
            name_path = osp.join(path,'names.txt')
            graph_lab_path = osp.join(path,'graph_labels.txt')
            region_path = osp.join(path,'region.txt')
            # Read data from `raw_path`.
            xl,yl,xh,yh = np.loadtxt(region_path)
            with open(graph_lab_path,'r') as f:
                rWLs = np.array([float(line) if line != '\n' else 0 for line in f.readlines()])
            with open(name_path,'r') as f:
                names = [int(line) if line != '\n' else 0 for line in f.readlines()]
            # normalize
            rWLs = rWLs/(xh-xl+yh-yl)
            rWLs = rWLs/1000.0

            mylist = np.append(mylist,rWLs[names])
        
        maxx = int(mylist.max())+1
        step = maxx/self.classes
        for i in range(0,self.tot_file_num):
            data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
            data.y = torch.tensor(int(data.wl.item() // step),dtype=torch.long)
            if data.y >= self.classes:
                data.y = self.classes - 1
            torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))

    def set_y_as_wl(self):
        for i in range(0,self.tot_file_num):
            data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
            data.y = torch.tensor(data.wl,dtype=torch.float)
            torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))




    def process(self):

        i = 0
        for design in self.raw_file_names:
            # paths
            path = osp.join(self.raw_dir,design)
            net_path = osp.join(path,'nets.txt')
            nodes_path = osp.join(path,'node_feature.txt')
            name_path = osp.join(path,'names.txt')
            pos_path = osp.join(path,'pos.txt')
            graph_lab_path = osp.join(path,'graph_labels.txt')
            pin_path = osp.join(path,'pins.txt')
            region_path = osp.join(path,'region.txt')
            golden_path = osp.join(path,'golden.txt')
            weight_path = osp.join(path,'node_pos.txt')
            # Read data from `raw_path`.
            golden = np.loadtxt(golden_path)
            nets = np.array(pd.read_table(net_path,header=None))
            node_lab = np.array(pd.read_table(nodes_path,header=None),dtype=float)
            pins_ = np.array(pd.read_table(pin_path,header=None),dtype=int)
            pin_feature = np.array(pd.read_table(pin_path,header=None),dtype=float)[:,2:]
            xl,yl,xh,yh = np.loadtxt(region_path)
            node_weight = np.loadtxt(weight_path)
            with open(name_path,'r') as f:
                names = [int(line) if line != '\n' else 0 for line in f.readlines()]
            pos = np.array(pd.read_table(pos_path,header=None),dtype=float)
            with open(graph_lab_path,'r') as f:
                rWLs = np.array([float(line) if line != '\n' else 0 for line in f.readlines()])

            # stastics
            #rel = (rWLs[names]/golden).tolist()
            #pylab.hist(rel,20)
            #pylab.xlabel('Range')
            #pylab.ylabel('Count')
            #pylab.savefig('{}.png'.format(design+'_golden'))
            #pylab.cla()
            #rWLs = np.array(pd.read_table(graph_lab_path,header=None),dtype=float)
            # get pos
            macro_indeces = torch.tensor(pos[0],dtype=int)
            macro_size = node_lab[macro_indeces]
            pos = pos[1:]
            # normalize
            node_lab[:,0] = node_lab[:,0]/(xh-xl)
            node_lab[:,1] = node_lab[:,1]/(yh-yl)
            pin_feature[:,0] = pin_feature[:,0]/(xh-xl)
            pin_feature[:,1] = pin_feature[:,1]/(yh-yl)
            # std
            rWLs = rWLs/(xh-xl+yh-yl)
            rWLs = rWLs/1000.0
            orWLs = rWLs.copy()
            mean = np.mean(rWLs[names])
            std = np.std(rWLs[names])
            #rWLs = np.tanh((rWLs - mean)/std)
            rWLs = np.arctan((rWLs - mean)/std)/1.58
        
            level = np.zeros((rWLs.shape[0]),dtype=int)
            step = 2.0/self.classes

            for q in range(0,self.classes):
                level += np.where((rWLs >= q*step -1.0) & (rWLs < (q+1)*step-1.0),q,0)
                #print(level.max())
                print(((rWLs >= q*step -1.0) & (rWLs < (q+1)*step-1.0))[names].sum())
                #level[:,q] = np.where((rWLs >= q*step -1.0) & (rWLs < (q+1)*step-1.0),1,0)

            #check = np.sum(level,axis=1)
            #print(check)
            cell_size = torch.tensor(node_lab, dtype=torch.float)
            edge_index = torch.tensor(nets.T, dtype=torch.long)
            pins = torch.tensor(pin_feature,dtype=torch.float)
            gold = torch.tensor(golden,dtype=torch.float)
            weight = torch.tensor(node_weight,dtype=torch.float)
            for name in names:
                my_pos = pos[2*name:2*name+2].T
                # normalize pos
                my_pos[:,0] = my_pos[:,0]-xl
                my_pos[:,1] = my_pos[:,1]-yl
                my_pos[:,0] = my_pos[:,0]/(xh-xl)
                my_pos[:,1] = my_pos[:,1]/(yh-yl)

                macro_pos = torch.tensor(my_pos,dtype=torch.float)
                # vx,_=torch.max(weight,dim=-1)
                # vn,_=torch.min(weight,dim=-1)
                # print((vx-vn).max())
                # print((vx-vn).min())
                # pylab.hist((vx-vn).cpu().numpy(),20)
                # pylab.xlabel('Range')
                # pylab.ylabel('Count')
                # pylab.savefig('{}.png'.format(design+'_golden'))
                # pdb.set_trace()
                cell_pos = torch.matmul(weight,macro_pos)
                cell_pos[macro_indeces] = macro_pos
                x = torch.cat([cell_size,cell_pos],dim=-1)
                y = torch.tensor(level[name],dtype=torch.long)
                if y.ndim <= 1:
                    y = y.view(1,-1)
                data = Data(x=x, edge_index=edge_index,y=y,pin_feature=pins,macro_pos=macro_pos,macro_index=macro_indeces,golden=gold,wl=orWLs[name])

                if self.pre_filter is not None and not self.pre_filter(data):
                    continue

                if self.pre_transform is not None:
                    data = self.pre_transform(data)

                torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
                i += 1

    def len(self):
        return len(self.processed_file_names)

    def get(self, idx):
        data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(idx)))
        return data


class CNNDataset(Dataset):
    def __init__(self, root, transform=None, pre_transform=None, classes=5):
        self.tot_file_num = None # int
        self.file_num = None # dict, file nums for each design
        self.ptr = None
        self.classes = classes
        self.root = root
        super(CNNDataset, self).__init__(root, transform, pre_transform)


    @property
    def processed_dir(self):
        return osp.join(self.root,'imgs')


    @property
    def raw_file_names(self):
        return ['des_perf_a', 'edit_dist_a', 'fft_a','fft_b','matrix_mult_a','matrix_mult_b','matrix_mult_c','pci_bridge32_b']


    @property
    def num_classes(self):
        return int(self.classes)

    @property
    def processed_file_names(self):
        if self.tot_file_num is None:
            self.tot_file_num = 0
            self.file_num = {}
            self.ptr = {}
            for design in self.raw_file_names:
                path = osp.join(self.raw_dir,design)
                name_path = osp.join(path,'names.txt')
                names = np.array(pd.read_table(name_path,header=None)).reshape(-1)
                self.tot_file_num += names.shape[0]
                self.file_num[design] = names.shape[0]
            self.ptr[self.raw_file_names[0]] = 0
            for i in range(1,int(len(self.raw_file_names))):
                self.ptr[self.raw_file_names[i]] = self.ptr[self.raw_file_names[i-1]] + self.file_num[self.raw_file_names[i-1]]
        return ['data_%d.pt'%i for i in range(0,self.tot_file_num)]


    def reset_y(self,classes=None):
        if classes is not None:
            self.classes = classes
        mylist = np.array([])
        for design in self.raw_file_names:
            # paths
            path = osp.join(self.raw_dir,design)
            name_path = osp.join(path,'names.txt')
            graph_lab_path = osp.join(path,'graph_labels.txt')
            region_path = osp.join(path,'region.txt')
            # Read data from `raw_path`.
            xl,yl,xh,yh = np.loadtxt(region_path)
            with open(graph_lab_path,'r') as f:
                rWLs = np.array([float(line) if line != '\n' else 0 for line in f.readlines()])
            with open(name_path,'r') as f:
                names = [int(line) if line != '\n' else 0 for line in f.readlines()]
            # normalize
            rWLs = rWLs/(xh-xl+yh-yl)
            rWLs = rWLs/1000.0

            mylist = np.append(mylist,rWLs[names])
        
        # stastics
        mylist = np.array(mylist).reshape(-1,)
        rel = mylist.tolist()
        pylab.hist(rel,20)
        pylab.xlabel('Range')
        pylab.ylabel('Count')
        pylab.savefig('{}.png'.format('Total'))
        pylab.cla()


        maxx = int(mylist.max())+1
        step = maxx/self.classes
        for i in range(0,self.tot_file_num):
            data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
            data.y = torch.tensor(int(data.wl.item() // step),dtype=torch.long)
            if data.y >= self.classes:
                data.y = self.classes - 1
            torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))


    def set_y_as_rank(self,classes=None):
        if classes is not None:
            self.classes = classes
        mylist = np.array([])
        for design in self.raw_file_names:
            # paths
            path = osp.join(self.raw_dir,design)
            name_path = osp.join(path,'names.txt')
            graph_lab_path = osp.join(path,'graph_labels.txt')
            region_path = osp.join(path,'region.txt')
            # Read data from `raw_path`.
            xl,yl,xh,yh = np.loadtxt(region_path)
            with open(graph_lab_path,'r') as f:
                rWLs = np.array([float(line) if line != '\n' else 0 for line in f.readlines()])
            with open(name_path,'r') as f:
                names = [int(line) if line != '\n' else 0 for line in f.readlines()]
            # normalize
            rWLs = rWLs/(xh-xl+yh-yl)
            rWLs = rWLs/1000.0

            mylist = np.append(mylist,rWLs[names])
        
        maxx = int(mylist.max())+1
        step = maxx/self.classes
        for i in range(0,self.tot_file_num):
            data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
            data.y = torch.tensor(int(data.wl.item() // step),dtype=torch.long)
            if data.y >= self.classes:
                data.y = self.classes - 1
            torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))

    def set_y_as_wl(self):
        for i in range(0,self.tot_file_num):
            data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
            data.y = torch.tensor(data.wl,dtype=torch.float)
            torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))



    def process(self):

        i = 0
        for design in self.raw_file_names:
            # paths
            path = osp.join(self.raw_dir,design)
            net_path = osp.join(path,'nets.txt')
            nodes_path = osp.join(path,'node_feature.txt')
            name_path = osp.join(path,'names.txt')
            pos_path = osp.join(path,'pos.txt')
            graph_lab_path = osp.join(path,'graph_labels.txt')
            pin_path = osp.join(path,'pins.txt')
            region_path = osp.join(path,'region.txt')
            golden_path = osp.join(path,'golden.txt')
            # Read data from `raw_path`.
            golden = np.loadtxt(golden_path)
            nets = np.array(pd.read_table(net_path,header=None))
            node_lab = np.array(pd.read_table(nodes_path,header=None),dtype=float)
            pins_ = np.array(pd.read_table(pin_path,header=None),dtype=int)
            pin_feature = np.array(pd.read_table(pin_path,header=None),dtype=float)[:,2:]
            xl,yl,xh,yh = np.loadtxt(region_path)
            with open(name_path,'r') as f:
                names = [int(line) if line != '\n' else 0 for line in f.readlines()]
            pos = np.array(pd.read_table(pos_path,header=None),dtype=float)
            with open(graph_lab_path,'r') as f:
                rWLs = np.array([float(line) if line != '\n' else 0 for line in f.readlines()])

            # stastics
            #rel = (rWLs[names]/golden).tolist()
            #pylab.hist(rel,20)
            #pylab.xlabel('Range')
            #pylab.ylabel('Count')
            #pylab.savefig('{}.png'.format(design+'_golden'))
            #pylab.cla()
            #rWLs = np.array(pd.read_table(graph_lab_path,header=None),dtype=float)
            # get pos
            macro_indeces = torch.tensor(pos[0],dtype=int)
            macro_size = node_lab[macro_indeces]
            pos = pos[1:]
            # normalize
            node_lab[:,0] = node_lab[:,0]/(xh-xl)
            node_lab[:,1] = node_lab[:,1]/(yh-yl)
            pin_feature[:,0] = pin_feature[:,0]/(xh-xl)
            pin_feature[:,1] = pin_feature[:,1]/(yh-yl)
            # std
            
            rWLs = rWLs/(xh-xl+yh-yl)
            rWLs = rWLs/1000.0
            orWLs = rWLs.copy()
            mean = np.mean(rWLs[names])
            std = np.std(rWLs[names])
            #rWLs = np.tanh((rWLs - mean)/std)
            rWLs = np.arctan((rWLs - mean)/std)/1.58
        
            level = np.zeros((rWLs.shape[0]),dtype=int)
            step = 2.0/self.classes

            for q in range(0,self.classes):
                level += np.where((rWLs >= q*step -1.0) & (rWLs < (q+1)*step-1.0),q,0)
                #print(level.max())
                #print(((rWLs >= q*step -1.0) & (rWLs < (q+1)*step-1.0))[names].sum())
                #level[:,q] = np.where((rWLs >= q*step -1.0) & (rWLs < (q+1)*step-1.0),1,0)

            #check = np.sum(level,axis=1)
            #print(check)
            x = torch.tensor(node_lab, dtype=torch.float)
            edge_index = torch.tensor(nets.T, dtype=torch.long)
            pins = torch.tensor(pin_feature,dtype=torch.float)
            gold = torch.tensor(golden,dtype=torch.float)

            kernel = GaussianBlur().cuda()

            for name in names:
                my_pos = pos[2*name:2*name+2].T

                pos_pic = x.new_zeros(int(xh-xl),int(yh-yl),dtype=torch.float)
                pin_pic = x.new_zeros(int(xh-xl),int(yh-yl),dtype=torch.float)
                net_pic = x.new_zeros(int(xh-xl),int(yh-yl),dtype=torch.float)
                one = x.new_ones(1,dtype=torch.float)
                for j in range(0,int(len(macro_indeces))):
                    ind = int(macro_indeces[j])
                    a= int(my_pos[j][0])
                    b= int(macro_size[j][0])
                    c= int(my_pos[j][1])
                    d= int(macro_size[j][1])
                    pos_pic[a:a+b,c:c+d] = 1.0
                    
                    pins = pins_[pins_[:,0] == ind]

                    padding = 10

                    # for pin in pins:
                    #     pin_pic[a+pin[2]-padding:a+pin[2]+padding,c+pin[3]-padding:c+pin[3]+padding] = 1.0
                    #     net_id = int(pin[1])
                    #     deg = len(pins_[pins_[:,1] == net_id]) + 64
                    #     if deg>255.0:
                    #         deg = 255.0
                    #     net_pic[a+pin[2]-padding:a+pin[2]+padding,c+pin[3]-padding:c+pin[3]+padding] += deg/255.0
                    
                    for pin in pins:
                        if a+pin[2] < pin_pic.shape[0] and c+pin[3] < pin_pic.shape[1]:
                            pin_pic[a+pin[2],c+pin[3]] = .5
                            net_id = int(pin[1])
                            deg = len(pins_[pins_[:,1] == net_id]) + 64
                            if deg>255.0:
                                deg = 255.0
                            net_pic[a+pin[2],c+pin[3]] += deg/255.0

                pos_pic=kernel(pos_pic.cuda())
                pin_pic=kernel(pin_pic.cuda())
                net_pic=kernel(net_pic.cuda())
                #net_pic = torch.clamp(net_pic,0,1)
                #im = Image.fromarray(np.uint8((net_pic*255).view(net_pic.shape[2],net_pic.shape[3]).numpy()))
                #im.convert('L')
                #im.save('test.jpg') 

                pos_pic = torch.nn.functional.interpolate(pos_pic,size=(224,224))
                pin_pic = torch.nn.functional.interpolate(pin_pic,size=(224,224))
                net_pic = torch.nn.functional.interpolate(net_pic,size=(224,224))

                pos_pic = torch.clamp(pos_pic,0,1)
                pin_pic = torch.clamp(pin_pic,0,1)
                net_pic = torch.clamp(net_pic,0,1)

                pic = torch.cat([pos_pic,pin_pic,net_pic],dim=1)


                y = torch.tensor(level[name],dtype=torch.long)
                if y.ndim <= 1:
                    y = y.view(1,-1)
                data = Data(x=torch.zeros(1),dge_index=None,y=y,picture=pic,wl=torch.tensor(orWLs[name],dtype=torch.float))

                if self.pre_filter is not None and not self.pre_filter(data):
                    continue

                if self.pre_transform is not None:
                    data = self.pre_transform(data)

                torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
                i += 1
                print(i)

    def len(self):
        return len(self.processed_file_names)

    def get(self, idx):
        transform = transforms.Compose([
            transforms.Resize(224),
            transforms.Normalize((.5,.5,.5),(.5,.5,.5))
            ])
        data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(idx)))
        data.picture = transform(data.picture)
        return data


class DesignSet(Dataset):
    def __init__(self, root, transform=None, pre_transform=None, classes=20):
        self.tot_file_num = None # int
        self.file_num = None # dict, file nums for each design
        self.ptr = None
        self.classes = classes
        self.label_key = None
        super(DesignSet, self).__init__(root, transform, pre_transform)

    @property
    def raw_file_names(self):
        names_path = osp.join(self.root,'raw','all.names')
        names = np.loadtxt(names_path,dtype=str)
        return names.tolist()
    
    @property
    def train_file_names(self):
        names_path = osp.join(self.root,'raw','train.names')
        names = np.loadtxt(names_path,dtype=str)
        return names.tolist()

    @property
    def test_file_names(self):
        names_path = osp.join(self.root,'raw','test.names')
        names = np.loadtxt(names_path,dtype=str)
        return names.tolist()

    @property
    def num_features(self):
        return self[0].x.size(1) + self[0].cell_pos.size(1) + self[0].pin_feature.size(1) + 1 # macro label

    @property
    def num_classes(self):
        return int(self.classes)

    @property
    def processed_file_names(self):
        if self.tot_file_num is None:
            self.tot_file_num = 0
            self.file_num = {}
            self.ptr = {}
            for design in self.raw_file_names:
                path = osp.join(self.raw_dir,design)
                name_path = osp.join(path,'names.txt')
                names = np.array(pd.read_table(name_path,header=None)).reshape(-1)
                self.tot_file_num += names.shape[0]
                self.file_num[design] = names.shape[0]
            self.ptr[self.raw_file_names[0]] = 0
            for i in range(1,int(len(self.raw_file_names))):
                self.ptr[self.raw_file_names[i]] = self.ptr[self.raw_file_names[i-1]] + self.file_num[self.raw_file_names[i-1]]
        return ['data_%d.pt'%i for i in range(0,self.tot_file_num)]


    def reset_wl(self,transform):
        for i in range(0,self.tot_file_num):
            data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
            data.wl = transform(data.wl)
            torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))


    def reset_rank(self,classes):
        if self.classes == classes:
            return 
        mylist = np.array([])
        for design in self.raw_file_names:
            # paths
            path = osp.join(self.raw_dir,design)
            name_path = osp.join(path,'names.txt')
            graph_lab_path = osp.join(path,'graph_labels.txt')
            region_path = osp.join(path,'region.txt')
            # Read data from `raw_path`.
            xl,yl,xh,yh = np.loadtxt(region_path)
            with open(graph_lab_path,'r') as f:
                rWLs = np.array([float(line) if line != '\n' else 0 for line in f.readlines()])
            with open(name_path,'r') as f:
                names = [int(line) if line != '\n' else 0 for line in f.readlines()]
            # normalize
            rWLs = rWLs/(xh-xl+yh-yl)
            rWLs = rWLs/1000.0

            mylist = np.append(mylist,rWLs[names])
        
        maxx = int(mylist.max())+1
        step = maxx/self.classes
        for i in range(0,self.tot_file_num):
            data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
            data.rank = torch.tensor(int(data.wl.item() // step),dtype=torch.long)
            if data.rank >= self.classes:
                data.rank = self.classes - 1
            torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))


    def process(self):

        i = 0
        for design in self.raw_file_names:
            # paths
            path = osp.join(self.raw_dir,design)
            size_path = osp.join(path,'node_size.txt')
            name_path = osp.join(path,'names.txt')
            pos_root = osp.join(path,'node_pos')
            wl_path = osp.join(path,'wl.txt')
            pin_path = osp.join(path,'pins.txt')
            region_path = osp.join(path,'region.txt')
            golden_path = osp.join(path,'golden.txt')
            dist_path = osp.join(path,'dist2macro.txt')
            macro_path = osp.join(path,'macro_index.txt')
            hpwl_path = osp.join(path,'hpwl.txt')
            # Read data from `raw_path`.
            golden = np.loadtxt(golden_path)
            pins = np.loadtxt(pin_path)
            size = np.loadtxt(size_path)
            incidence = pins[:,:2]
            pin_feature = pins[:,2:]
            xl,yl,xh,yh = np.loadtxt(region_path)
            dist2macro = np.loadtxt(dist_path)
            macro_index = torch.tensor(np.loadtxt(macro_path),dtype=torch.long)
            names = np.loadtxt(name_path,dtype=int)
            with open(wl_path,'r') as f:
                rWLs = np.array([float(line) if line != '\n' else 0 for line in f.readlines()])
            
            with open(hpwl_path,'r') as f:
                hpwls = np.array([float(line) if line != '\n' else 0 for line in f.readlines()])

            # stastics
            rel = (rWLs[names]).tolist()
            pylab.hist(rel,20)
            pylab.xlabel('Range')
            pylab.ylabel('Count')
            pylab.savefig('stat/{}.png'.format(design+'_golden'))
            pylab.cla()
            # normalize
            size[:,0] = size[:,0]/(xh-xl)
            size[:,1] = size[:,1]/(yh-yl)
            pin_feature[:,0] = pin_feature[:,0]/(xh-xl)
            pin_feature[:,1] = pin_feature[:,1]/(yh-yl)
            # std
            rWLs = rWLs/(xh-xl+yh-yl)
            rWLs = rWLs/1000.0

            hpwls = hpwls/(xh-xl+yh-yl)
            hpwls = hpwls/1000.0

            orWLs = rWLs.copy()
            mean = np.mean(rWLs[names])
            std = np.std(rWLs[names])
            #rWLs = np.tanh((rWLs - mean)/std)
            rWLs = np.arctan((rWLs - mean)/std)/1.58
        
            level = np.zeros((rWLs.shape[0]),dtype=int)
            step = 2.0/self.classes

            for q in range(0,self.classes):
                level += np.where((rWLs >= q*step -1.0) & (rWLs < (q+1)*step-1.0),q,0)
                #print(((rWLs >= q*step -1.0) & (rWLs < (q+1)*step-1.0))[names].sum())

            cell_size = torch.tensor(size, dtype=torch.float)
            edge_index = torch.tensor(incidence.T, dtype=torch.long)
            pins = torch.tensor(pin_feature,dtype=torch.float)
            gold = torch.tensor(golden,dtype=torch.float)
            weight = torch.tensor(dist2macro,dtype=torch.float)
            weight = 1/weight
            summ = weight.sum(dim=-1).view(-1,1).repeat(1,int(len(macro_index)))
            weight = weight/summ
            #weight = torch.softmax(1/weight,dim=-1)
            for name in names:
                if hpwls[name] == 0:
                    print('{}-{}'.format(design,name))
                pos_path = osp.join(pos_root,'%d.txt'%name)
                node_pos = torch.tensor(np.loadtxt(pos_path),dtype=torch.float) 

                node_pos[:,0] = (node_pos[:,0]-xl)/(xh-xl)
                node_pos[:,1] = (node_pos[:,1]-yl)/(yh-yl)

                macro_pos = node_pos[macro_index]

                fake_pos = torch.matmul(weight,macro_pos)
                fake_pos[macro_index] = macro_pos
        
                y = torch.tensor(level[name],dtype=torch.long)
                if y.ndim <= 1:
                    y = y.view(1,-1)
                data = Data(x=cell_size[:node_pos.shape[0]], 
                            edge_index=edge_index,
                            y=y,
                            pin_feature=pins,
                            macro_index=macro_index,
                            golden=gold,
                            wl=torch.tensor(orWLs[name],dtype=torch.float),
                            rank=y,
                            cell_pos=node_pos,
                            fake_pos=fake_pos,
                            hpwl=torch.tensor(hpwls[name],dtype=torch.float))

                if self.pre_filter is not None and not self.pre_filter(data):
                    continue

                if self.pre_transform is not None:
                    data = self.pre_transform(data)

                torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
                i += 1

    def len(self):
        return len(self.processed_file_names)

    def get(self, idx):
        data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(idx)))
        return data


Set = DesignSet('data')
print(Set.raw_file_names)


